'''
文件说明
用途：对yolo检测结果格式进行转换
yolo检测结果（txt文件格式）：class, x_center, y_center, width, height, conf
注意：yolo检测结果为缩放后百分比（相对原图）

转换后txt文件格式：x_top_left, y_top_left, width, height, conf


author：周小龙
date：24-12-1
'''
from pathlib import Path
import numpy as np

"""
对yolo检测格式进行转换
yolo->erti
"""
# 图片默认宽和高
width=1280
height=1024

def yolo2erti(in_path:Path, out_path:Path, width=1280, height=1024):
    '''
    对yolo检测结果进行转换：yolo -> erti
    转换后格式为：x_top_left, y_top_left, width, height, conf

    :param in_path: yolo检测结果label文件夹（txt文件目录）
    :param out_path: 输出文件目录
    :param width:
    :param height:
    :return:
    '''
    files=[label for label in in_path.rglob('*') if label.is_file() and label.suffix in ['.txt']]
    if not out_path.exists():
        # 目标目录不存在
        out_path.mkdir()

    for file in files:
        labels=np.loadtxt(file)
        if labels.ndim==1:
            new_labels=np.zeros((1,5))
            new_labels[:,0] =np.round((labels[1]-labels[3]/2)*width,4)
            new_labels[:,1] =np.round((labels[2]-labels[4]/2)*height,4)
            new_labels[:,2]=np.round(labels[3]*width,4) # width
            new_labels[:,3]=np.round(labels[4]*height,4) # height
            new_labels[:,4]=np.round(labels[5],4) # confidence
        else:
            row=len(labels)
            new_labels=np.zeros((row,5))
            new_labels[:,0]=np.round((labels[:,1]-labels[:,3]/2)*width,4)
            new_labels[:,1]=np.round((labels[:,2]-labels[:,4]/2)*height,4)
            new_labels[:,2]=np.round(labels[:,3]*width,4) # width
            new_labels[:,3]=np.round(labels[:,4]*height,4) # height
            new_labels[:,4]=np.round(labels[:,5],4) # confidence

        np.savetxt(Path(out_path,file.name),new_labels,fmt='%0.4f',delimiter=",")
    print('转换完成')



